Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@cpersson-amd I've been out on PTO for a month. I'll take a closer look at this next week. Meanwhile, can you update your branch with the latest in main. Thanks.

@susanbao can you take a quick look at this PR.

@cpersson-amd please review Sanbao's comments above and rebase with main. We tested the PR internally and it looks good. Would you be willing to change the axis fsdp to context? If not, I can make the change after this PR is merged.

@entrpn I've rebased with main, included @susanbao requested change and updated the mesh names: fsdp -> context, fsdp_batch -> fsdp. Please let me know if anything else needs to be changed.

@entrpn I've rebased with main, included @susanbao requested change and updated the mesh names: fsdp -> context, fsdp_batch -> fsdp. Please let me know if anything else needs to be changed.

thanks @cpersson-amd this looks great. Can you run ruff check --fix as the unit tests are failing due to formatting right now.

@entrpn Sure, I ran 'ruff check --fix' and had to manually fix some bare except statements. It should be good with the latest commit

@cpersson-amd Please review my PR to fix some of the unit tests. Once they pass, this can be merged. cpersson-amd#1